package security;

import java.io.BufferedReader;
import java.io.FileInputStream;
import java.io.FileNotFoundException;
import java.io.FileReader;
import java.io.IOException;
import java.io.InputStreamReader;
import java.net.Socket;
import java.security.InvalidKeyException;
import java.security.KeyPair;
import java.security.NoSuchAlgorithmException;
import java.security.PrivateKey;
import java.security.SecureRandom;
import java.security.Key;

import javax.crypto.KeyGenerator;
import javax.crypto.SecretKey;
import javax.crypto.spec.IvParameterSpec;
import javax.crypto.Mac;
import javax.crypto.spec.SecretKeySpec;

import org.bouncycastle.openssl.PEMReader;
import org.bouncycastle.openssl.PasswordFinder;

import org.bouncycastle.util.encoders.Base64;
import org.bouncycastle.util.encoders.Hex;

import security.channel.AESChannel;
import security.channel.Channel;
import security.channel.RSAChannel;
import security.channel.TCPChannel;

public class SecureHelper {
	public static PrivateKey getPrivateKey(String file) throws IOException {
		PEMReader pemIn = null;
		pemIn = new PEMReader(new FileReader(file), new PasswordFinder() {

			@Override
			public char[] getPassword() {
				// reads the password from standard input for decrypting
				// the private key
				String pass = "";
				System.out.println("Enter pass phrase:");
				try {
					pass = new BufferedReader(new InputStreamReader(System.in))
							.readLine();
				} catch (IOException e) {
					System.err.println("ERROR reading password: "
							+ e.getMessage());
				}
				return pass.toCharArray();
			}

		});

		KeyPair keyPair = (KeyPair) pemIn.readObject();
		pemIn.close();

		return keyPair.getPrivate();
	}

	public static AESChannel createAESChannelClient(String keyServerPub,
			String keyDirClient, String loginCommand, int tcpPort, Socket socket) {

		String username = loginCommand.split(" ")[1].toLowerCase();
		RSAChannel rsaChannel = null;
		AESChannel aesChannel = null;
		PrivateKey privateKey = null;

		// load privateKey of the user
		try {
			privateKey = SecureHelper.getPrivateKey(keyDirClient + username
					+ ".pem");
		} catch (FileNotFoundException e) {
			System.err.println("You have no private key!");
			return null;
		} catch (IOException e) {
			System.err.println("Wrong password!");
			return null;
		}

		try {
			rsaChannel = new RSAChannel(socket, keyServerPub, privateKey);
		} catch (IOException e) {
			System.err.println("Couldn't create RSAChannel for login client: "
					+ e.getMessage());
		}

		// create !login-message
		String challengeClient = createSecureRandom(32);
		loginCommand += " " + tcpPort + " " + challengeClient;

		rsaChannel.println(loginCommand);

		String message2 = null;
		try {
			message2 = rsaChannel.readLine();
		} catch (IOException e) {
			System.out.println("ERROR, receiving 2. message from server.");
		}

		// check if public-key of client was not found on server
		if (message2.equals("!notok no pub key")) {
			System.out
					.println("Your public key file doesn't exist on the server!");
			return null;
		}

		byte[] keyAESBytes = Base64.decode(message2.split(" ")[3]);
		SecretKey keyAES = new SecretKeySpec(keyAESBytes, 0,
				keyAESBytes.length, "AES");

		// Initialisation vector:
		byte[] iv = Base64.decode(message2.split(" ")[4]);
		IvParameterSpec ivParameterSpec = new IvParameterSpec(iv);

		try {
			aesChannel = new AESChannel(socket, keyAES, ivParameterSpec);
		} catch (IOException e) {
			System.err.println("ERROR: creating aesChannel in SecureHelper: "
					+ e.getMessage());
		}
		
		// check server sent correct challange
		if (!challengeClient.equals(message2.split(" ")[1])) {
			aesChannel.println("!notok clientchallange wrong");
			System.err.println("Server sent wrong clientchallange back!");
			return null;
		}

		aesChannel.println(message2.split(" ")[2]); // send server-challange

		String serverReceive = "";
		try {
			serverReceive = aesChannel.readLine();
			System.out.println(serverReceive);
		} catch (IOException e) {
			System.out.println("ERROR, receiving 4. message from server");
		}

		if (serverReceive.startsWith("User already in use") || serverReceive.startsWith("!notok")) {
			return null;
		}

		return aesChannel;
	}

	public static AESChannel createAESChannelServer(String keyDirClient,
			String inputLine, Socket socket) {
		String username = inputLine.split(" ")[1].toLowerCase();
		Channel channel = null;
		AESChannel aesChannel = null;
		
		

		try {
			channel = new RSAChannel(socket, keyDirClient + username
					+ ".pub.pem", null);
		} catch (IOException e) { // wenn public key doesn't exist
			try {
				channel = new TCPChannel(socket);
			} catch (IOException e1) {
				System.err
						.println("Error creating TCPChannel in SecureHelper: "
								+ e1.getMessage());
			}
			channel.println("!notok no pub key"); // send !notok and return null
			return null;
		}

		// create 2. message of the handshake.
		String challengeServer = SecureHelper.createSecureRandom(32);
		String okMessage = "!ok " + inputLine.split(" ")[3] + " "
				+ challengeServer;

		KeyGenerator generator = null;
		try {
			generator = KeyGenerator.getInstance("AES");
		} catch (NoSuchAlgorithmException e) {
			System.err.println("ERROR: wrong algorithm creating AES-key: "
					+ e.getMessage());
		}
		generator.init(256);
		SecretKey keyAES = generator.generateKey();
		byte[] keyAESBytes = keyAES.getEncoded();
		String keyAESString = new String(Base64.encode(keyAESBytes));

		// Initialisation vector:
		String ivparamString = SecureHelper.createSecureRandom(16);
		byte[] iv = Base64.decode(ivparamString);
		IvParameterSpec ivParameterSpec = new IvParameterSpec(iv);

		okMessage += " " + keyAESString + " " + ivparamString;

		channel.println(okMessage);

		try {
			aesChannel = new AESChannel(socket, keyAES, ivParameterSpec);
		} catch (IOException e) {
			System.err.println("Error creating AESCHannel in SecureHelper: "
					+ e.getMessage());
		}
		channel = aesChannel;

		String message3 = "";
		try {
			message3 = channel.readLine();
		} catch (IOException e) {
			System.err.println("Error receiving 3. message of handshake!: "
					+ e.getMessage());
		}
		
		if(message3.equals("!notok clientchallange wrong")) {
			System.err.println("Client said wrong clientchallange was sent back! Login denied.");
			return null;
		}

		if (!challengeServer.equals(message3)) { //if client sent wrong challange back
			try {
				channel = new TCPChannel(socket);
			} catch (IOException e1) {
				System.err
						.println("Error creating TCPChannel in SecureHelper: "
								+ e1.getMessage());
			}
			
			channel.println("!notok serverchallenge wrong");			
			System.err.println("Client sent wrong challange back! Login denied.");
			return null;
		}
		return aesChannel;
	}

	/**
	 * Creates a secure random number and encode it with base64.
	 * 
	 * @param bytes
	 *            the count of bytes used for the secred number.
	 * @return the secure random number.
	 */
	public static String createSecureRandom(int bytes) {
		SecureRandom secureRandom = new SecureRandom();
		final byte[] challengeBytes = new byte[bytes];
		secureRandom.nextBytes(challengeBytes);
		byte[] challengeBytes64 = Base64.encode(challengeBytes);
		return new String(challengeBytes64);
	}

	/**
	 * Creates a shared secret key from a file, used for creating HMACs
	 * 
	 * @param pathToSecretKey
	 *            The path to the keyfile
	 * @return The shared secret key
	 * @throws IOException
	 *             thrown if file was not found or errors reading or closing the
	 *             file
	 */
	public static Key createSharedSecretKey(String pathToSecretKey)
			throws IOException {
		byte[] keyBytes = new byte[1024];
		FileInputStream fis = new FileInputStream(pathToSecretKey);
		fis.read(keyBytes);
		fis.close();
		byte[] input = Hex.decode(keyBytes);
		Key key = new SecretKeySpec(input, "HmacSHA256");
		return key;
	}

	/**
	 * Creates a hash MAC for a String.
	 * 
	 * @param secretKey
	 *            The shared secret key used to generate the hash.
	 * @param message
	 *            The message the hash should be created from.
	 * @return The hash MAC.
	 * @throws InvalidKeyException
	 *             thrown if key is invalid
	 */
	public static byte[] createHashMAC(Key secretKey, String message)
			throws InvalidKeyException {
		Mac hMac = null;
		try {
			hMac = Mac.getInstance("HmacSHA256");
			hMac.init(secretKey);
			// MESSAGE is the message to sign in bytes
			hMac.update(message.getBytes());
			byte[] hash = hMac.doFinal();
			return hash;
		} catch (NoSuchAlgorithmException e) {
			e.printStackTrace();
		}
		return null;
	}
}
